
import tensorflow as tf
import os
import numpy as np
from PIL import Image
import sys
sys.path.append('/home/dl/offline/jesselv.comp.jpg/quiz-w7-code')
from datasets import dataset_utils


tf.flags.DEFINE_string('train_dataset_dir', './train', '')
tf.flags.DEFINE_float('validation_rate', 0.15, '')
tf.flags.DEFINE_string('output_dir', './tfrecord', '')
tf.flags.DEFINE_string('train_output', 'train.tfrecord', '')
tf.flags.DEFINE_string('validation_output', 'validation.tfrecord', '')

#tf.flags.DEFINE_string('train_dataset_dir', './', '')
#tf.flags.DEFINE_float('validation_rate', 0.3, '')
#tf.flags.DEFINE_string('output_dir', './', '')
#tf.flags.DEFINE_string('train_output', 'train.tfrecord', '')
#tf.flags.DEFINE_string('validation_output', 'validation.tfrecord', '')

FLAGS = tf.app.flags.FLAGS


def _extract_image(filename):
    img = Image.open(filename)
    if img is None:
        raise ValueError('can not open {0}'.format(filename))

    height = img.size[1]
    width = img.size[0]

    with open(filename, 'rb') as f:
        buf = f.read()
    return buf, height, width


def _add_to_tfrecord(data, height, width, label, tfrecord_writer, encoded_jpg, sess):
    jpg_string = sess.run(encoded_jpg, feed_dict={'input_image:0': data})
    example = dataset_utils.image_to_tfexample(
        jpg_string, 'jpg'.encode(), height, width, label)
    tfrecord_writer.write(example.SerializeToString())


def convert_to_tfrecord():
    if not tf.gfile.IsDirectory(FLAGS.train_dataset_dir):
        raise ValueError('train_dataset_dir must be a dir')
    if FLAGS.validation_rate <= 0 or FLAGS.validation_rate >= 1:
        raise ValueError('validation_rate value must be in range 0 to 1')
    if not tf.gfile.IsDirectory(FLAGS.output_dir):
        raise ValueError('output_dir must be a dir')
    if not isinstance(FLAGS.train_output, str) or FLAGS.train_output == '':
        raise ValueError('train_output must be a instance of str, and can not be a empty str')
    if not isinstance(FLAGS.validation_output, str) or FLAGS.validation_output == '':
        raise ValueError('validation_output must be a instance of str, and can not be a empty str')

    graph = tf.Graph()
    with graph.as_default():
        image = tf.placeholder(dtype=tf.uint8, name='input_image')
        encoded_jpg = tf.image.encode_jpeg(image)

        image_encode = tf.placeholder(dtype=tf.string, name='input_encode')
        decode_data = tf.image.decode_jpeg(image_encode)

    handwriting_dir_list = tf.gfile.ListDirectory(FLAGS.train_dataset_dir)

    label_map = {}

    train_output_writer \
        = tf.python_io.TFRecordWriter(os.path.join(FLAGS.output_dir, FLAGS.train_output))
    validation_output_writer = tf.python_io.TFRecordWriter(os.path.join(FLAGS.output_dir, FLAGS.validation_output))

    len_word_size = len(handwriting_dir_list)

    validation_size = len_word_size * 400 * FLAGS.validation_rate
    validation_size = int(validation_size)
    if validation_size <= 0:
        validation_size = 1

    shuffle_data = np.arange(len_word_size*400)
    np.random.shuffle(shuffle_data)
    validation_data = shuffle_data[: validation_size]
    train_data = shuffle_data[validation_size:]

    validation_cnt = 0
    train_cnt = 0
    with tf.Session(graph=graph) as sess:

        word_dict = {}
        label_dict = {}

        for i, dir in enumerate(handwriting_dir_list):
            if not os.path.isdir(os.path.join(FLAGS.train_dataset_dir, dir)):
                raise ValueError('{0} is not a dir'.format(dir))

            label_map[i] = dir
            word_list = tf.gfile.ListDirectory(os.path.join(FLAGS.train_dataset_dir, dir))

            for j, image_name in enumerate(word_list):
                image_file = os.path.join(FLAGS.train_dataset_dir, dir, image_name)
                if not os.path.isfile(image_file):
                    raise ValueError('{0} is not a file'.format(image_file))

                word_dict[i*400+j] = image_file
                label_dict[i*400+j] = i

        for i in validation_data:
            image_file = word_dict[i]
            data, height, width = _extract_image(image_file)
            data = sess.run(decode_data, feed_dict={'input_encode:0': data})
            _add_to_tfrecord(data, height, width, label_dict[i], validation_output_writer, encoded_jpg, sess)
            validation_cnt += 1
            #print('{0} --> {1}'.format(image_file.split('/')[-2], label_map.get(label_dict[i])))
            if validation_cnt % 100 == 0:
                print('creating validation tfrecord {0}'.format(validation_cnt))

        for i in train_data:
            image_file = word_dict[i]
            data, height, width = _extract_image(image_file)
            data = sess.run(decode_data, feed_dict={'input_encode:0': data})
            _add_to_tfrecord(data, height, width, label_dict[i], train_output_writer, encoded_jpg, sess)
            train_cnt += 1
            #print('{0} --> {1}'.format(image_file.split('/')[-2], label_map.get(label_dict[i])))
            if train_cnt % 100 == 0:
                print('creating train tfrecord {0}'.format(train_cnt))

        dataset_utils.write_label_file(label_map, FLAGS.output_dir)
        validation_output_writer.close()
        train_output_writer.close()

        print('validation image cnt {0}, train image cnt {1}'.format(validation_cnt, train_cnt))


convert_to_tfrecord()
